from mpi4py import MPI
import numpy as np
from petsc4py import PETSc
import ufl
from ufl import dx, ds, dS, grad, inner, dot, avg, jump
from dolfinx import mesh, fem, io, la
import basix
from dolfinx.fem.petsc import assemble_matrix, assemble_vector


from mpi4py import MPI
import numpy as np
from petsc4py import PETSc
import ufl
from ufl import dx, ds, dS, grad, inner, dot, avg, jump
from dolfinx import mesh, fem
from dolfinx.fem.petsc import assemble_matrix, assemble_vector, apply_lifting, set_bc
from dolfinx import io

# ------------------ 函数定义 ------------------
def u_exact_func(x, t):
    return np.exp(x[0] + x[1] + t)

def u_exact_expr(x, t):
    return ufl.exp(x[0] + x[1] + t)

def f_func(x, t):
    return -np.exp(x[0] + x[1] + t)

# ------------------ 参数 ------------------
type_ = "SIPG"            
epsilon = -1.0       
k = 2                     
sigma_0_base = 6
x_min, x_max, y_min, y_max = 0.0, 2.0, 0.0, 1.0
T = 1.0

h_list = [1/2, 1/4, 1/8, 1/16, 1/32]
h1_semi_err, l2_err = [], []

for h in h_list:
    dt = h**2
    N = int(round(1 / h))
    domain = mesh.create_rectangle(
        MPI.COMM_WORLD,
        [[x_min, y_min], [x_max, y_max]],
        [int((x_max - x_min) / h), int((y_max - y_min) / h)],
        mesh.CellType.triangle,
    )

    x = ufl.SpatialCoordinate(domain)
    V = fem.functionspace(domain, ("DG", k))

    u = ufl.TrialFunction(V)
    v = ufl.TestFunction(V)

    n = ufl.FacetNormal(domain)
    h_cell = ufl.CellDiameter(domain)
    h_avg = ufl.avg(h_cell)

    sigma_int = sigma_0_base * (k + 1) * k / h_avg
    sigma_bnd = sigma_0_base * (k + 1) * k / h_cell

    # ------- DG离散化矩阵 A -------
    a_vol = inner(grad(u), grad(v)) * dx
    a_int = (
        - inner(avg(grad(v)), jump(u, n)) * dS
        + epsilon * inner(avg(grad(u)), jump(v, n)) * dS
        + sigma_int * inner(jump(u, n), jump(v, n)) * dS
    )
    a_D = (
        epsilon * dot(grad(u), n) * v * ds
        - dot(grad(v), n) * u * ds
        + sigma_bnd * u * v * ds
    )
    A_form = a_vol + a_int + a_D

    # 质量矩阵
    M_form = inner(u, v) * dx

    A = assemble_matrix(fem.form(A_form))
    A.assemble()
    M = assemble_matrix(fem.form(M_form))
    M.assemble()

    # 初值
    u_n = fem.Function(V)
    u_n.interpolate(lambda x: np.exp(x[0] + x[1] + 0.0))

    u_new = fem.Function(V)

    # 右端项形式
    f_expr = fem.Function(V)


    t = 0.0
    A_half = M.copy()
    A_half.axpy(dt/2, A, False)  # m + dt/2 * A
    A_half.assemble()

    kspCN = PETSc.KSP().create(MPI.COMM_WORLD)
    kspCN.setOperators(A_half)
    kspCN.setType("preonly")
    kspCN.getPC().setType("lu")

    # crank-nicolson处理时间步
    L_vec_old = None
    while t < T - 1e-10:
        t += dt
        # 当前右端体项 + 边界项
        f_expr.interpolate(lambda x: f_func(x, t))
        L_form = inner(f_expr, v) * dx \
            + epsilon * dot(grad(v), n) * u_exact_expr(x, t) * ds \
            + sigma_bnd * u_exact_expr(x, t) * v * ds
        L_vec = assemble_vector(fem.form(L_form))
        L_vec.assemble()

        # 第一拍用 BE 启动（L_vec_old 不存在）
        if L_vec_old is None:
            L_vec_old = L_vec.copy()

        # CN 右端  b = (M - 0.5*dt*A) u_n  +  0.5*dt*(L^n + L^{n+1})
        b = M.createVecRight()
        M.mult(u_n.x.petsc_vec, b)                       # b = M u_n
        # createVecRight等价于：先拿一个“空壳”向量，形状跟 A 的行对齐，然后让矩阵把乘积填进去。
        # 它本身不隐含“右端项”意思，只是“给我一块能装 A*x 的内存”。
        Au_tmp = A.createVecRight()
        A.mult(u_n.x.petsc_vec, Au_tmp)
        b.axpy(-dt/2, Au_tmp)                            # b -= 0.5*dt*A u_n
        b.axpy(dt/2, L_vec_old)                          # b += 0.5*dt*L^n
        b.axpy(dt/2, L_vec)                              # b += 0.5*dt*L^{n+1}

        kspCN.solve(b, u_new.x.petsc_vec)
        u_new.x.scatter_forward()

        # 前进
        u_n.x.array[:] = u_new.x.array[:]
        L_vec_old = L_vec.copy()
    
    

    # -------- 误差计算 --------
    u_e = fem.Function(V)
    u_e.interpolate(lambda x: u_exact_func(x, T))

    diff = fem.Function(V)
    diff.x.array[:] = u_n.x.array[:] - u_e.x.array[:]

    L2_err = np.sqrt(fem.assemble_scalar(fem.form(inner(diff, diff) * dx)))
    H1_semi = np.sqrt(fem.assemble_scalar(fem.form(inner(grad(diff), grad(diff)) * dx)))

    l2_err.append(L2_err)
    h1_semi_err.append(H1_semi)

    if MPI.COMM_WORLD.rank == 0:
        print(f"h={h:.4f}, L2误差={L2_err:.3e}, H1半范数误差={H1_semi:.3e}")

# -------- 收敛阶统计 --------
def _rates(errs, hs):
    out = []
    for i in range(len(errs) - 1):
        out.append(np.log(errs[i] / errs[i+1]) / np.log(hs[i] / hs[i+1]))
    return out

l2_list = _rates(l2_err, h_list)
h1_semi_list = _rates(h1_semi_err, h_list)

if MPI.COMM_WORLD.rank == 0:
    print("\n===== 收敛结果 =====")
    print("h:               ", ["{:.5f}".format(h) for h in h_list])
    print("L2收敛阶:         ", ["{:.2f}".format(e) for e in l2_list])
    print("H1半范数收敛阶:   ", ["{:.2f}".format(e) for e in h1_semi_list])